import torch
import wandb
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import random

from torchmetrics import Metric
from lightning import Callback
from io import BytesIO
from omegaconf import DictConfig
from PIL import Image


class PlotECGCallback(Callback):
    def __init__(self, n_samples=1, leads_to_plot=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], lead_order=None, log_every_n_epochs=50):
        super().__init__()
        self.n_samples = n_samples
        self.leads_to_plot = leads_to_plot
        self.lead_order = lead_order if lead_order is not None else ["I", "II", "III", "AVR", "AVL", "AVF", "V1", "V2", "V3", "V4", "V5", "V6"]
        self.log_every_n_epochs = log_every_n_epochs
        self.last_batch = None

    def format_ecg_plot(self, graph, data, label, v_min, v_max):
        """ Adds ECG-style grid lines and labels to a plot. """
        graph.set_ylim([v_min, v_max])
        sns.despine()
        graph.grid(True, which='both', linestyle='-', linewidth=0.5, color='black', alpha=0.6)
        graph.axhline(0, color='red', linewidth=1.5, linestyle='--')  # Add prominent zero line
        graph.set_ylabel(label, fontsize=12)
        graph.minorticks_on()
        graph.tick_params(which='both', top=True, left=True, right=False, bottom=True)
        
        yticks_up = np.arange(0, v_max + 1.0, 1.0)
        yticks_down = np.arange(0, v_min - 1.0, -1.0)
        yticks = np.concatenate([yticks_down[::-1][:-1], yticks_up])  # Combine up and down ticks
        graph.set_yticks(yticks)
        graph.set_yticklabels([f'{y:.1f}' for y in yticks], fontsize=10)
        
        xticks = np.arange(0, len(data) + 1, 250) 
        graph.set_xticks(xticks)
        graph.set_xticklabels([f'{x}' for x in xticks], fontsize=10)

    def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0):
        if trainer.current_epoch % self.log_every_n_epochs == 0:
            self.last_batch = {
                'original': outputs['original'].detach().cpu(),
                'masked': outputs['masked'].detach().cpu(),
                'reconstructed': outputs['reconstructed'].detach().cpu(),
                'im_paste': outputs['im_paste'].detach().cpu()
            }

    def on_validation_epoch_end(self, trainer, pl_module):
        if trainer.current_epoch % self.log_every_n_epochs == 0:
            if self.last_batch is None:
                return

            original = self.last_batch['original'].to(torch.float32).numpy()
            masked = self.last_batch['masked'].to(torch.float32).numpy()
            reconstructed = self.last_batch['reconstructed'].to(torch.float32).numpy()
            im_paste = self.last_batch['im_paste'].to(torch.float32).numpy()

            batch_size = original.shape[0]
            sample_count = min(self.n_samples, batch_size)
            indices = np.random.choice(batch_size, size=sample_count, replace=False)

            for idx in indices:
                log_images = {}
                for lead_idx in self.leads_to_plot:
                    lead_name = self.lead_order[lead_idx]
                    fig, axes = plt.subplots(4, 1, figsize=(12, 12), constrained_layout=True)
                    fig.suptitle(f'Sample {idx + 1} - Lead {lead_name}', fontsize=16)

                    # Determine common y-axis limits
                    v_min = min(np.min(original[idx, 0, lead_idx]), np.min(masked[idx, 0, lead_idx]), np.min(reconstructed[idx, 0, lead_idx]))
                    v_max = max(np.max(original[idx, 0, lead_idx]), np.max(masked[idx, 0, lead_idx]), np.max(reconstructed[idx, 0, lead_idx]))
                    v_max = max(abs(v_max), abs(v_min)) * 1.1  # Add 10% padding for visibility
                    v_min = -v_max

                    # Original
                    self.format_ecg_plot(axes[0], original[idx, 0, lead_idx], 'Original', v_min, v_max)
                    sns.lineplot(x=np.arange(original.shape[-1]), y=original[idx, 0, lead_idx], ax=axes[0], color='black', linewidth=1.0)

                    # Masked
                    self.format_ecg_plot(axes[1], masked[idx, 0, lead_idx], 'Masked', v_min, v_max)
                    sns.lineplot(x=np.arange(masked.shape[-1]), y=masked[idx, 0, lead_idx], ax=axes[1], color='black', linewidth=1.0)

                    # Reconstructed
                    self.format_ecg_plot(axes[2], reconstructed[idx, 0, lead_idx], 'Reconstructed', v_min, v_max)
                    sns.lineplot(x=np.arange(reconstructed.shape[-1]), y=reconstructed[idx, 0, lead_idx], ax=axes[2], color='black', linewidth=1.0)
                    
                    # MAE Reconstruction pasted with visible patches
                    self.format_ecg_plot(axes[3], im_paste[idx, 0, lead_idx], 'Reconstructed + Visible Patches', v_min, v_max)
                    sns.lineplot(x=np.arange(im_paste.shape[-1]), y=im_paste[idx, 0, lead_idx], ax=axes[3], color='black', linewidth=1.0)

                    log_images[f"ECG Reconstruction Lead {lead_name}"] = wandb.Image(fig)
                    plt.close(fig)

                log_images['trainer/global_step'] = trainer.global_step
                
                trainer.logger.experiment.log(log_images)

            self.last_batch = None


class ECGLoggerMetricPyTorch(Metric):
    def __init__(self, cfg: DictConfig,
                 ecg_model,
                 lead_order: list = ["I", "II", "III", "AVR", "AVL", "AVF", "V1", "V2", "V3", "V4", "V5", "V6"],
                 leads_to_plot : list = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
                 num_samples: int = 1
                 ):
        super().__init__()

        self.cfg = cfg
        self.ecg_model = ecg_model
        self.num_samples = num_samples
        self.lead_order = lead_order
        self.leads_to_plot = leads_to_plot
        self.batch = None

    def update(self, outputs, batch):
        logits = outputs.logits
        mask = outputs.mask
        
        mask = mask.unsqueeze(-1).repeat(1, 1, self.ecg_model.config.patch_size[0] * self.ecg_model.config.patch_size[1])
        mask = self.ecg_model.unpatchify(mask)
        reconstructed = self.ecg_model.unpatchify(logits)
        
        masked_pixel_values = batch * (1 - mask)
        ecg_paste = batch * (1 - mask) + reconstructed * mask

        self.batch = {
            'original': batch.detach().cpu().to(torch.float32).squeeze(),
            'masked': masked_pixel_values.detach().cpu().to(torch.float32).squeeze(),
            'reconstructed': reconstructed.detach().cpu().to(torch.float32).squeeze(),
            'ecg_paste': ecg_paste.detach().cpu().to(torch.float32).squeeze(),
        }

    def compute(self):
        selected_index = random.randint(0, self.batch['original'].size(0) - 1)
        selected_batch = {kind: self.batch[kind][selected_index] for kind in self.batch.keys()}
        return selected_batch
    
    def format_ecg_plot(self, graph, data, label, v_min, v_max):
        """ Adds ECG-style grid lines and labels to a plot. """
        graph.set_ylim([v_min, v_max])
        sns.despine()
        graph.grid(True, which='both', linestyle='-', linewidth=0.5, color='black', alpha=0.6)
        graph.axhline(0, color='red', linewidth=1.5, linestyle='--')  # Add prominent zero line
        graph.set_ylabel(label, fontsize=12)
        graph.minorticks_on()
        graph.tick_params(which='both', top=True, left=True, right=False, bottom=True)
        
        yticks_up = np.arange(0, v_max + 1.0, 1.0)
        yticks_down = np.arange(0, v_min - 1.0, -1.0)
        yticks = np.concatenate([yticks_down[::-1][:-1], yticks_up])  # Combine up and down ticks
        graph.set_yticks(yticks)
        graph.set_yticklabels([f'{y:.1f}' for y in yticks], fontsize=10)
        
        xticks = np.arange(0, len(data) + 1, 250) 
        graph.set_xticks(xticks)
        graph.set_xticklabels([f'{x}' for x in xticks], fontsize=10)

    def log_plots(self):
        selected_batch = self.compute()

        for lead_idx in self.leads_to_plot:
            lead_name = self.lead_order[lead_idx]
            fig, axes = plt.subplots(4, 1, figsize=(12, 12), constrained_layout=True)
            fig.suptitle(f'ECG Lead {lead_name}', fontsize=16)

            original = selected_batch['original'].numpy()
            masked = selected_batch['masked'].numpy()
            reconstructed = selected_batch['reconstructed'].numpy()
            ecg_paste = selected_batch['ecg_paste'].numpy()

            # Determine common y-axis limits for better comparison
            v_min = min(np.min(original[lead_idx]), np.min(masked[lead_idx]), np.min(reconstructed[lead_idx]))
            v_max = max(np.max(original[lead_idx]), np.max(masked[lead_idx]), np.max(reconstructed[lead_idx]))
            v_max = max(abs(v_max), abs(v_min)) * 1.1  # Add 10% padding
            v_min = -v_max

            # Original ECG
            self.format_ecg_plot(axes[0], original[lead_idx], 'Original', v_min, v_max)
            sns.lineplot(x=np.arange(original.shape[-1]), y=original[lead_idx], ax=axes[0], color='black', linewidth=1.0)

            # Masked ECG
            self.format_ecg_plot(axes[1], masked[lead_idx], 'Masked', v_min, v_max)
            sns.lineplot(x=np.arange(masked.shape[-1]), y=masked[lead_idx], ax=axes[1], color='black', linewidth=1.0)

            # Reconstructed ECG
            self.format_ecg_plot(axes[2], reconstructed[lead_idx], 'Reconstructed', v_min, v_max)
            sns.lineplot(x=np.arange(reconstructed.shape[-1]), y=reconstructed[lead_idx], ax=axes[2], color='black', linewidth=1.0)
            
            # Reconstructed with visible patches
            self.format_ecg_plot(axes[3], ecg_paste[lead_idx], 'Reconstructed + Visible Patches', v_min, v_max)
            sns.lineplot(x=np.arange(ecg_paste.shape[-1]), y=ecg_paste[lead_idx], ax=axes[3], color='black', linewidth=1.0)

            # Log to Wandb
            wandb.log({f"ECG Reconstruction Lead {lead_name}": wandb.Image(fig)})
            plt.close(fig)

        self.reset()

    def reset(self):
        self.batch = None
